[レポート] Run high-performing FMs at scale with Amazon SageMaker HyperPod #AIM314 #AWSreInvent

AWS re:Invent 2024

#AWS

たかくに

2024.12.03

こんにちは！AWS 事業本部コンサルティング部のたかくに（@takakuni_）です。

re:Invent 2024 でラスベガスに来ています。

https://reinvent.awsevents.com/

「Run high-performing FMs at scale with Amazon SageMaker HyperPod」という、面白そうなワークショップがあったので参加してみました。

Amazon SageMaker HyperPod、名前を聞くだけで興奮しますね。

セッション概要

タイトル

AIM314 | Run high-performing FMs at scale with Amazon SageMaker HyperPod

説明

This workshop explores the end-to-end workflow of managing foundation models (FMs) on Amazon SageMaker HyperPod. The discussion covers four critical aspects: Parameter-Efficient Fine-Tuning (PEFT), model deployment, and serving. Additionally, dive into operational aspects, including system observability and the resiliency features of Amazon SageMaker HyperPod, such as recovery from failure and job auto-resuming. Leave this hands-on session with a robust understanding of managing and deploying foundation models efficiently on AWS. Learn to use cutting-edge techniques and tools to ensure high performance, reliability, and scalability in FM development. You must bring your laptop to participate.

スピーカー

Keita Watanabe, Sr. Solutions Architect, AWS
Alex Iankoulski, Principal Solutions Architect, AWS

その他サポートメンバーとして各国から多くの方が参加されてました。途中詰まったのですが、私の不慣れな英語でも懸命にサポートいただきありがとうございました。

内容

アジェンダ

アジェンダは以下の通りです。内容としては Amazon EKS Support in Amazon SageMaker HyperPod のショートカットバージョンです。

1. Cluster Setup
- a. Setup Environment Variables
- b. Configure the EKS Cluster
- c. Install Dependencies
- d. Create the HyperPod Cluster
- e. View the AWS Console
- f. Setup FSx for Lustre File System
1. Observability/Profiling
- 1. Amazon CloudWatch Container Insights
  - a. Container Insights setup
1. PyTorch DDP on CPU (t3/m5/c5)
1. Ray on HyperPod (g5,m5)
1. Observability/Profiling

ロングバージョンもあります。自分のアカウントでも試せるワークショップです。ぜひお試しください。

Cluster Setup

今回はオーケストレーターに EKS を利用した SageMaker Hyperpod を利用するため、 EKS クラスターのセットアップを行います。

ハンズオンでは EKS クラスターはすでに作成済みで、 Helm を使って HyperPod に必要なセットアップを行いました。

Install Dependencies

git clone https://github.com/aws/sagemaker-hyperpod-cli.git
cd sagemaker-hyperpod-cli/helm_chart

helm dependencies update HyperPodHelmChart
helm install hyperpod-dependencies HyperPodHelmChart

helm list

インストールしたパッケージは次のとおりです

Health Monitoring Agent
NVIDIA device plugin for Kubernetes
Neuron device plugin
EFA Kubernetes device plugin
Kubeflow Training Operator
Kubeflow MPI Operator
Kubernetes PriorityClass

Create the HyperPod Cluster

今回は GPU を使わず、 ml.m5.2xlarge で代用しました。

SageMaker HyperPod を作成するための JSON ファイルを作成します。Orchestrator には既存で作成していた EKS クラスターを紐付けます。

cluster-config.json

{
    "ClusterName": "ml-cluster",
    "Orchestrator": {
      "Eks":
      {
        "ClusterArn": "${EKS_CLUSTER_ARN}"
      }
    },
    "InstanceGroups": [
      {
        "InstanceGroupName": "worker-group-1",
        "InstanceType": "${ACCEL_INSTANCE_TYPE}",
        "InstanceCount": ${ACCEL_COUNT},
        "InstanceStorageConfigs": [
          {
            "EbsVolumeConfig": {
              "VolumeSizeInGB": ${ACCEL_VOLUME_SIZE}
            }
          }
        ],
        "LifeCycleConfig": {
          "SourceS3Uri": "s3://${BUCKET_NAME}",
          "OnCreate": "on_create.sh"
        },
        "ExecutionRole": "${EXECUTION_ROLE}",
        "ThreadsPerCore": 2
      }
    ],
    "VpcConfig": {
      "SecurityGroupIds": ["$SECURITY_GROUP"],
      "Subnets":["$SUBNET_ID"]
    },
    "NodeRecovery": "${NODE_RECOVERY}"
}

HyperPod クラスターを作成します。

aws sagemaker create-cluster \
    --cli-input-json file://cluster-config.json \
    --region $AWS_REGION
aws sagemaker list-clusters \
 --output table \
 --region $AWS_REGION

View the AWS Console

SageMaker HyperPod の良いところの 1 つはコンソールがあるところです。SageMaker コンソールに先ほど作成した HyperPod クラスターが表示されていますね。

2024-12-02 at 09.35.48-Amazon SageMaker us-west-2@2x.png

Setup FSx for Lustre File System

後続の DDP で利用する Lustre ファイルシステムを作成します。

sagemaker-user@default:~$ kubectl describe pvc fsx-claim
Name:          fsx-claim
Namespace:     default
StorageClass:  fsx-sc
Status:        Pending
Volume:
Labels:        <none>
Annotations:   volume.beta.kubernetes.io/storage-provisioner: fsx.csi.aws.com
               volume.kubernetes.io/storage-provisioner: fsx.csi.aws.com
Finalizers:    [kubernetes.io/pvc-protection]
Capacity:
Access Modes:
VolumeMode:    Filesystem
Used By:       <none>
Events:
  Type    Reason                Age              From                                                                                      Message
  ----    ------                ----             ----                                                                                      -------
  Normal  Provisioning          5s               fsx.csi.aws.com_fsx-csi-controller-5b9f6cf799-p974k_8f23a90e-6cd6-4074-a568-2f80d94a5355  External provisioner is provisioning volume for claim "default/fsx-claim"
  Normal  ExternalProvisioning  4s (x2 over 5s)  persistentvolume-controller                                                               Waiting for a volume to be created either by the external provisioner 'fsx.csi.aws.com' or manually by the system administrator. If volume creation is delayed, please verify that the provisioner is running and correctly registered.

Amazon CloudWatch Container Insights

Observability/Profiling ということで、Container Insights のセットアップを行います。

Hyperpod の実行ロールに CloudWatch Agent のポリシーをアタッチします。

export EX_ROLE_NAME=$(echo $EXECUTION_ROLE | sed 's/.*\///')
aws iam attach-role-policy \
--role-name $EX_ROLE_NAME \
--policy-arn arn:aws:iam::aws:policy/CloudWatchAgentServerPolicy

EKS のアドオンに amazon-cloudwatch-observability を追加します。

aws eks create-addon --addon-name amazon-cloudwatch-observability --cluster-name $EKS_CLUSTER_NAME

うまくエージェントが動いていますね。

sagemaker-user@default:~$ kubectl get pods -n amazon-cloudwatch
NAME                                                              READY   STATUS    RESTARTS   AGE
amazon-cloudwatch-observability-controller-manager-57fcb55gtdcb   1/1     Running   0          5m57s
cloudwatch-agent-hvsks                                            1/1     Running   0          5m54s
cloudwatch-agent-lznhm                                            1/1     Running   0          5m54s
fluent-bit-ptqfj                                                  1/1     Running   0          5m57s
fluent-bit-q8j8x                                                  1/1     Running   0          5m57s

PyTorch DDP on CPU (t3/m5/c5)

Setup

今回は SageMaker Hyperpod で PyTorch を使った CPU (ml.m5.2xlarge) インスタンスでの分散データ並列（DDP）を行いました。

まずは Docker イメージの作成から ECR へのプッシュまでです。ファイルシステムには先ほど作成した FSx Lustre を利用します。

cd ~
git clone https://github.com/aws-samples/awsome-distributed-training/
cd awsome-distributed-training/3.test_cases/16.pytorch-cpu-ddp/kubernetes

export AWS_REGION=$(aws ec2 describe-availability-zones --output text --query 'AvailabilityZones[0].[RegionName]')
export ACCOUNT=$(aws sts get-caller-identity --query Account --output text)
export REGISTRY=${ACCOUNT}.dkr.ecr.${AWS_REGION}.amazonaws.com/
docker build $DOCKER_NETWORK -t ${REGISTRY}fsdp:pytorch2.2-cpu ..

# Create registry if needed
REGISTRY_COUNT=$(aws ecr describe-repositories | grep \"fsdp\" | wc -l)
if [ "$REGISTRY_COUNT" == "0" ]; then
        aws ecr create-repository --repository-name fsdp
fi

# Login to registry
echo "Logging in to $REGISTRY ..."
aws ecr get-login-password | docker login --username AWS --password-stdin $REGISTRY

# Push image to registry
docker image push ${REGISTRY}fsdp:pytorch2.2-cpu

Training

モデルのトレーニングに移ります。

今回用に PVC の名前とエポック数を変更します。

fsdp.yaml

apiVersion: v1
kind: Service
metadata:
  name: etcd
spec:
  ports:
    - name: etcd-client-port
      port: 2379
      protocol: TCP
      targetPort: 2379
  selector:
    app: etcd

---
apiVersion: apps/v1
kind: Deployment
metadata:
  labels:
    app: etcd
  name: etcd
spec:
  replicas: 1
  selector:
    matchLabels:
      app: etcd
  template:
    metadata:
      labels:
        app: etcd
    spec:
      containers:
        - name: etcd
          command: ["/usr/local/bin/etcd"]
          args:
            - "--data-dir"
            - "/var/lib/etcd"
            - "--enable-v2"
            - "--listen-client-urls"
            - "http://0.0.0.0:2379"
            - "--advertise-client-urls"
            - "http://0.0.0.0:2379"
            - "--initial-cluster-state"
            - "new"
          image: quay.io/coreos/etcd:latest
          ports:
            - containerPort: 2379
              name: client
              protocol: TCP
            - containerPort: 2380
              name: server
              protocol: TCP
      restartPolicy: Always
---
apiVersion: "kubeflow.org/v1"
kind: PyTorchJob
metadata:
  name: fsdp
spec:
  elasticPolicy:
    rdzvBackend: etcd
    rdzvHost: etcd
    rdzvPort: 2379
    minReplicas: 1
    maxReplicas: 64
    maxRestarts: 100
    metrics:
      - type: Resource
        resource:
          name: cpu
          target:
            type: Utilization
            averageUtilization: 90
  pytorchReplicaSpecs:
    Worker:
      replicas: 2
      restartPolicy: OnFailure
      template:
        metadata:
          labels:
            app: fsdp
        spec:
          volumes:
            - name: shmem
              hostPath:
                path: /dev/shm
            - name: local
              hostPath:
                path: /mnt/k8s-disks/0
            - name: fsx-pv
              persistentVolumeClaim:
+                claimName: fsx-claim
+                # claimName: fsx-pvc
          #nodeSelector:
          #  node.kubernetes.io/instance-type: "ml.m5.2xlarge"
          containers:
            - name: pytorch
              image: 509677657506.dkr.ecr.us-west-2.amazonaws.com/fsdp:pytorch2.2-cpu
              imagePullPolicy: Always
              command:
                - /opt/conda/bin/torchrun
                - --nproc_per_node=4
                - --nnodes=2
                - /workspace/ddp.py
+                - "50000"
                - "10"
                - --batch_size=32
                - --checkpoint_path=/fsx/snapshot.pt

              volumeMounts:
                - name: shmem
                  mountPath: /dev/shm
                - name: local
                  mountPath: /local
                - name: fsx-pv
                  mountPath: /fsx

fsdp.yaml-template の内容をアプライします。

kubectl apply -f ./fsdp.yaml

うまく動いていますね。

sagemaker-user@default:~$ kubectl get po
NAME                                                         READY   STATUS    RESTARTS   AGE
etcd-6bcc78d587-dltq9                                        1/1     Running   0          5m32s
fsdp-worker-0                                                1/1     Running   0          5m32s
fsdp-worker-1                                                1/1     Running   0          5m32s
hyperpod-dependencies-hyperpod-helm-chart-6f8989f9bb-54bpn   1/1     Running   0          58m
hyperpod-dependencies-mpi-operator-574c8c7f-bbpzh            1/1     Running   0          58m

次のステップように削除しちゃいます。

kubectl delete -f ./fsdp.yaml

Ray on HyperPod (g5,m5)

Ray は並列な Python アプリケーションを実行するように設計されたオープンソースの分散コンピューティングフレームワークです。私も今回はじめて触りました。

Ray クラスタ、ジョブ、サービスデプロイと管理に必要なツールをコンテナ化した aws-do-ray リポジトリを利用しました。

sagemaker-user@default:~/awsome-distributed-training/3.test_cases/16.pytorch-cpu-ddp/kubernetes$ cd ~

mkdir -p ~/wd

cp env_vars wd

docker run -it $DOCKER_NETWORK -v ${HOME}/wd:/wd -v ${HOME}/.aws:/root/.aws -v ${HOME}/.kube:/root/.kube --workdir /ray public.ecr.aws/hpc-cloud/aws-do-ray bash
Unable to find image 'public.ecr.aws/hpc-cloud/aws-do-ray:latest' locally
latest: Pulling from hpc-cloud/aws-do-ray
7478e0ac0f23: Already exists
a3a939b9441c: Pull complete
e627462bdfb5: Pull complete
8f04ec8bab35: Pull complete
Digest: sha256:ad8cc8699911f39ed4af4117fdf227f8719a0ff8530ffb7466fd0f70072a8a9a
Status: Downloaded newer image for public.ecr.aws/hpc-cloud/aws-do-ray:latest
                                 __
  ____ __      _______      ____/ /___        _________ ___  __
 / __ `/ | /| / / ___/_____/ __  / __ \______/ ___/ __ `/ / / /
/ /_/ /| |/ |/ (__  )_____/ /_/ / /_/ /_____/ /  / /_/ / /_/ /
\__,_/ |__/|__/____/      \__,_/\____/     /_/   \__,_/\__, /
                                                      /____/

aws-do-ray shell v20241125

alias - show list of command shortcuts
ll - list files in current directory

total 12K
drwxr-xr-x 7 root root  140 Nov 26 07:58 .
drwxr-xr-x 1 root root   28 Dec  2 17:59 ..
drwxr-xr-x 7 root root   99 Nov 26 07:58 deploy
drwxr-xr-x 3 root root 4.0K Nov 26 07:58 ops
drwxr-xr-x 3 root root  231 Nov 26 07:58 raycluster
drwxr-xr-x 3 root root  142 Nov 26 07:58 rayjob
drwxr-xr-x 6 root root  245 Nov 26 07:58 rayservice
-rwxrwxr-x 1 root root  615 Nov 26 07:58 remove-dependencies.sh
-rwxrwxr-x 1 root root 1.3K Nov 26 07:58 setup-dependencies.sh

それではデプロイします。

# Create KubeRay namespace
kubectl create namespace kuberay
# Deploy the KubeRay operator with the Helm chart repository
helm repo add kuberay https://ray-project.github.io/kuberay-helm/
helm repo update
#Install both CRDs and Kuberay operator v1.1.0
helm install kuberay-operator kuberay/kuberay-operator --version 1.1.0 --namespace kuberay
# Kuberay operator pod will be deployed onto head pod
kubectl get pods --namespace kuberay

うまくデプロイできていますね。Workshop は時間の関係上ここでフィニッシュでした。

root@default:/ray/rayservice#   kubectl get pods -n kuberay
NAME                                READY   STATUS    RESTARTS   AGE
kuberay-operator-69d46987ff-r6k7r   1/1     Running   0          16m